import os
import numpy as np
import torch
from tensorboardX import SummaryWriter
from utils.shared_buffer import SharedReplayBuffer
from algorithms.mat.mat_trainer import MATTrainer as TrainAlgo
from algorithms.mat.algorithm.transformer_policy import TransformerPolicy as Policy


# 将PyTorch张量转换为NumPy数组
def _t2n(x):
    return x.detach().cpu().numpy()


class Runner(object):
    """
    多智能体强化学习训练的基础运行器类
    负责管理训练循环、数据收集、模型更新和评估等核心功能
    :param config: (dict) 包含训练参数的配置字典
    """

    def __init__(self, config):

        # 从配置中提取基本参数
        self.all_args = config['all_args']                                          # 所有命令行参数
        self.envs = config['envs']                                                  # 训练环境
        self.eval_envs = config['eval_envs']                                        # 评估环境
        self.device = config['device']                                              # 计算设备（CPU/GPU）
        self.num_agents = config['num_agents']
        self.n_rollout_threads = self.all_args.n_rollout_threads
        self.n_eval_rollout_threads = self.all_args.n_eval_rollout_threads

        # 训练参数配置
        self.num_env_steps = self.all_args.num_env_steps                            # 总的环境步数
        self.episode_length = self.all_args.episode_length                          # 每个episode的长度

        # 训练间隔配置
        self.save_interval = self.all_args.save_interval  # 模型保存间隔
        self.use_eval = self.all_args.use_eval  # 是否进行评估
        self.eval_interval = self.all_args.eval_interval  # 评估间隔
        self.log_interval = self.all_args.log_interval  # 日志记录间隔

        # 模型目录配置
        self.model_dir = self.all_args.model_dir  # 预训练模型目录

        # 设置实验日志和模型保存目录
        # 手动创建目录
        self.run_dir = config["run_dir"]
        self.log_dir = str(self.run_dir / 'logs')  # 日志目录
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)
        self.writter = SummaryWriter(self.log_dir)  # TensorBoard写入器
        self.save_dir = str(self.run_dir / 'models')  # 模型保存目录
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)

        # 打印环境配置信息
        # 兼容不同的环境包装器类型
        if hasattr(self.envs, 'envs'):
            # DummyVecEnv类型
            print("obs_shape: ", self.envs.envs[0].obs_shape)
            print("act_shape: ", self.envs.envs[0].act_shape)
            obs_shape = self.envs.envs[0].obs_shape
            act_shape = self.envs.envs[0].act_shape
        else:
            # ShareDummyVecEnv类型
            print("obs_shape: ", self.envs.obs_shape)
            print("act_shape: ", self.envs.act_shape)
            obs_shape = self.envs.obs_shape
            act_shape = self.envs.act_shape
        
        self.policy = Policy(self.all_args,
                             obs_shape,  # 观测形状
                             act_shape,  # 动作数量
                             self.num_agents,  # 智能体数量
                             device=self.device)  # 计算设备

        # 如果指定了预训练模型，则加载模型
        if self.model_dir is not None:
            self.restore(self.model_dir)

        # 创建训练器（算法）
        self.trainer = TrainAlgo(self.all_args, self.policy, self.num_agents, device=self.device)

        # 创建经验回放缓冲区
        self.buffer = SharedReplayBuffer(self.all_args,
                                         self.num_agents,
                                         obs_shape,  # 使用处理后的观测空间
                                         act_shape)  # 使用处理后的动作空间

    # 主训练循环：收集训练数据、执行训练更新、评估策略
    def run(self):
        raise NotImplementedError

    # 收集预训练的热身数据、在正式训练前收集一些初始数据来初始化缓冲区
    def warmup(self):
        raise NotImplementedError

    # 收集训练用的轨迹数据
    def collect(self, step):
        raise NotImplementedError

    # 将收集的数据插入到训练缓冲区
    def insert(self, data):
        raise NotImplementedError

    # 计算价值函数和回报、用于PPO等算法中的优势函数计算
    @torch.no_grad()  # 禁用梯度计算，提高效率
    def compute(self):
        self.trainer.prep_rollout()  # 准备策略网络为推理模式
        
        # 获取多智能体联合观测
        obs_batch = self.buffer.obs[-1]  # 形状: (n_rollout_threads, num_agents, 36)
        available_actions_batch = self.buffer.available_actions[-1]  # 形状: (n_rollout_threads, num_agents, 7)
        
        # 判断经验回放缓冲区中是否有"可用动作掩码"
        # Critic网络在估值时会考虑当前哪些动作是可选的，避免对不可用动作做无意义的价值估计
        if self.buffer.available_actions is None:
            # 没有可用动作掩码时，只使用观测计算价值
            next_values = self.trainer.policy.get_values(obs_batch)
        else:
            # 有可用动作掩码时，同时考虑观测和可用动作
            next_values = self.trainer.policy.get_values(obs_batch, available_actions_batch)
        
        # 将价值函数输出分割为每个环境的独立数组
        next_values = np.array(np.split(_t2n(next_values), 1))
        
        # 使用计算出的价值函数更新缓冲区中的回报
        self.buffer.compute_returns(next_values, self.trainer.value_normalizer)

    # 使用缓冲区中的数据训练策略网络，并返回训练统计信息
    def train(self):

        self.trainer.prep_training()  # 准备策略网络为训练模式
        train_infos = self.trainer.train(self.buffer)  # 执行训练更新
        self.buffer.after_update()  # 训练后更新缓冲区状态
        return train_infos

    # 保存模型参数
    def save(self, episode):
        self.policy.save(self.save_dir, episode)

    # 恢复模型参数
    def restore(self, model_dir):
        self.policy.restore(model_dir)

    # 记录训练日志 :param train_infos: 训练信息字典 :param total_num_steps: 总训练步数
    def log_train(self, train_infos, total_num_steps):
        for k, v in train_infos.items():
            # 使用TensorBoard记录日志
            self.writter.add_scalars(k, {k: v}, total_num_steps)

    # 记录环境信息日志 :param env_infos: 环境信息字典 :param total_num_steps: 总训练步数
    def log_env(self, env_infos, total_num_steps):

        for k, v in env_infos.items():
            if len(v) > 0:  # 只记录非空的环境信息
                self.writter.add_scalars(k, {k: np.mean(v)}, total_num_steps)